import json
import os
# """
# create table electric_vehicle_sales
# (
#     id          int auto_increment primary key,
#     series_id   int            not null,
#     series_name varchar(255)   null,
#     brand_id    int            null,
#     brand_name  varchar(255)   null,
#     min_price   decimal(10, 2) null,
#     max_price   decimal(10, 2) null,
#     count       int            null,
#     time_period varchar(64)    not null
# );
#
# """
from sqlalchemy import create_engine, Column, Integer, String, DECIMAL, Enum
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

# 定义数据库模型
Base = declarative_base()

class ElectricVehicleSales(Base):
    __tablename__ = 'electric_vehicle_sales'
    id = Column(Integer, primary_key=True, autoincrement=True)  # 新增的自增主键
    series_id = Column(Integer)
    series_name = Column(String(255))
    brand_id = Column(Integer)
    brand_name = Column(String(255))
    min_price = Column(DECIMAL(10, 2))
    max_price = Column(DECIMAL(10, 2))
    count = Column(Integer)
    image = Column(String(255))
    time_period = Column(Enum('近一年', '近两年', '近三年'))  # 新增字段区分时间段

# 数据库连接配置
DATABASE_URI = 'mysql+pymysql://root:root@localhost/car_sales'

# 创建数据库引擎
engine = create_engine(DATABASE_URI)
Base.metadata.create_all(engine)  # 如果表不存在，则创建表

# 创建会话
Session = sessionmaker(bind=engine)
session = Session()

# 定义文件名和对应的时间段
# 请将下面的目录路径替换为你的实际目录路径
directory_path = './data/salesData'
files = [
    'car_sales_data_近一年.json',
    'car_sales_data_近两年.json',
    'car_sales_data_近三年.json'
]

# 遍历文件并存入数据库
for file_name in files:
    # 根据文件名提取时间周期
    if '近一年' in file_name:
        time_period = '近一年'
    elif '近两年' in file_name:
        time_period = '近两年'
    elif '近三年' in file_name:
        time_period = '近三年'
    else:
        print(f"无效的文件名: {file_name}")
        continue

    file_path = os.path.join(directory_path, file_name)  # 拼接文件路径

    # 检查文件是否存在
    if not os.path.exists(file_path):
        print(f"文件 {file_name} 不存在，跳过。")
        continue

    # 读取 JSON 文件
    with open(file_path, 'r', encoding='utf-8') as file:
        data = json.load(file)

    # 清洗和存入数据库
    for item in data['纯电动']:
        try:
            # 提取数据
            series_id = item['series_id']
            series_name = item['series_name']
            brand_id = item['brand_id']
            brand_name = item['brand_name']
            min_price = item['min_price']
            max_price = item['max_price']
            count = item['count']
            image = item['image']

            # 创建对象并添加到会话
            record = ElectricVehicleSales(
                series_id=series_id,
                series_name=series_name,
                brand_id=brand_id,
                brand_name=brand_name,
                min_price=min_price,
                max_price=max_price,
                count=count,
                time_period=time_period,
                image=image
            )
            session.add(record)
        except KeyError as e:
            print(f"KeyError: {e} in item: {item} (文件: {file_name})")
        except Exception as e:
            print(f"Error processing item: {item} (文件: {file_name}), Error: {e}")

# 提交事务
try:
    session.commit()
    print("数据已成功存入数据库！")
except Exception as e:
    session.rollback()
    print(f"事务提交失败: {e}")
finally:
    session.close()